import numpy as np
import torch
import torchvision
from matplotlib import pyplot as plt
from PIL import Image, ImageDraw
from depthaware.data.datasets_SSD_nyu import nyuSSDDataset
from depthaware.options.train_options import TrainOptions
from utils.utils import *
from tqdm import tqdm
from pprint import PrettyPrinter

from torchvision import transforms

# Good formatting when printing the APs for each class and mAP
pp = PrettyPrinter()

# Parameters
data_folder = './'
keep_difficult = True  # difficult ground truth objects must always be considered in mAP calculation, because these objects DO exist!
batch_size = 1
workers = 4
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
checkpoint = './checkpoint_ssd300.pth.tar'

# Load model checkpoint that is to be evaluated


# Load test data



def evaluate(test_loader, model):
    """
    Evaluate.

    :param test_loader: DataLoader for test data
    :param model: model
    """

    # Make sure it's in eval mode
    model.eval()

    # Lists to store detected and true boxes, labels, scores
    det_boxes = list()
    det_labels = list()
    det_scores = list()
    true_boxes = list()
    true_labels = list()
    true_difficulties = list()  # it is necessary to know which objects are 'difficult', see 'calculate_mAP' in utils.py

    # with torch.no_grad():
        # Batches
        # for i, (images, boxes, labels, _, difficulties) in enumerate(tqdm(test_loader, desc='Evaluating')):
    for i, (images, boxes, labels, depths, difficulties) in enumerate(tqdm(test_loader, desc='Evaluating')):
        images = images.to(device)  # (N, 3, 300, 300)
        depths = depths.to(device)

        # Forward prop.
        predicted_locs, predicted_scores = model(images, depths)

        # Detect objects in SSD output
        det_boxes_batch, det_labels_batch, det_scores_batch = model.detect_objects(predicted_locs, predicted_scores,
                                                                                   min_score=0.2, max_overlap=0.45,
                                                                                   top_k=200)
        # Evaluation MUST be at min_score=0.01, max_overlap=0.45, top_k=200 for fair comparision with the paper's results and other repos

        # Store this batch's results for mAP calculation
        boxes = [b.to(device) for b in boxes]
        labels = [l.to(device) for l in labels]
        difficulties = [d.to(device) for d in difficulties]

        det_boxes.extend(det_boxes_batch)
        det_labels.extend(det_labels_batch)
        det_scores.extend(det_scores_batch)
        true_boxes.extend(boxes)
        true_labels.extend(labels)
        true_difficulties.extend(difficulties)

        det_bo = det_boxes[0].cpu().detach().numpy().tolist()
        bo = boxes[0].cpu().detach().numpy().tolist()


        toPIL = transforms.ToPILImage()
        img = toPIL(images[0].cpu())
        # img = np.transpose(img, (1, 2, 0))
        # ima = Image.fromarray(img)
        draw = ImageDraw.Draw(img)  # 在上面画画
        for i in det_bo:
            i = [a * 300 for a in i]
            draw.rectangle(i, outline=(0, 0, 255))
        for i in bo:
            i = [a * 300 for a in i]
            draw.rectangle(i, outline=(255, 0, 0))
        img.show()

    pass
    # Calculate mAP
    # APs, mAP = calculate_mAP(det_boxes, det_labels, det_scores, true_boxes, true_labels, true_difficulties)
    #
    # # Print AP for each class
    # pp.pprint(APs)
    #
    # print('\nMean Average Precision (mAP): %.3f' % mAP)


if __name__ == '__main__':
    opt = TrainOptions().parse()
    test_dataset = nyuSSDDataset(opt, split='test')
    test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=False,
                                              collate_fn=test_dataset.collate_fn, num_workers=workers, pin_memory=True)

    checkpoint = torch.load(checkpoint)
    model = checkpoint['model']
    model = model.to(device)

    # Switch to eval mode
    model.eval()

    evaluate(test_loader, model)
